import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
import logging

def setup_logger():
    logger = logging.getLogger('mfl_framework')
    logger.setLevel(logging.INFO)
    
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    logger.addHandler(console_handler)
    
    file_handler = logging.FileHandler('mfl_experiment.log')
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)
    
    return logger

def load_datasets(num_clients, non_iid_degree):
    """
    模拟加载Non-IID数据集
    实际实现中应替换为真实数据集加载逻辑
    """
    train_datasets = []
    test_datasets = []
    
    # 模拟数据生成
    for i in range(num_clients):
        # 根据non_iid_degree创建有偏数据分布
        bias = non_iid_degree * i / num_clients
        train_data = _generate_synthetic_data(1000, bias)
        test_data = _generate_synthetic_data(200, bias)
        
        train_datasets.append(DataLoader(train_data, batch_size=32, shuffle=True))
        test_datasets.append(DataLoader(test_data, batch_size=32))
    
    return train_datasets, test_datasets

def _generate_synthetic_data(num_samples, bias):
    """生成合成数据集（实际应替换为真实DDoS数据集）"""
    class SyntheticDataset(Dataset):
        def __init__(self, num_samples, bias):
            self.data = torch.randn(num_samples, 1, 100)
            # 添加基于bias的标签偏移
            self.labels = (torch.rand(num_samples) > (0.5 + bias)).long()
        
        def __len__(self):
            return len(self.data)
        
        def __getitem__(self, idx):
            return self.data[idx], self.labels[idx]
    
    return SyntheticDataset(num_samples, bias)
